package com.ljx.rag.rag.service.impl;

import com.ljx.rag.rag.domain.dto.OllamaChatReq;
import com.ljx.rag.rag.domain.dto.OllamaChatResponse;
import com.ljx.rag.rag.domain.dto.OllamaMessage;
import com.ljx.rag.rag.service.AiService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.List;

/**
 * @Author: LiJunXi
 * @Description:
 * @Date: 2025-03-09  16:02
 */
@Service
public class AiServiceImpl implements AiService {

    @Autowired
    private WebClient webClient;

    @Value("${rag.ai.chat.model}")
    private String model;

    @Override
    public Flux<String> callChatApi(List<OllamaMessage> messages) {
        OllamaChatReq ollamaChatReq = OllamaChatReq.builder()
                .model(model)
                .messages(messages)
                .stream(true)
                .build();

        return webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.TEXT_EVENT_STREAM)
                .bodyValue(ollamaChatReq)
                .retrieve()
                .bodyToFlux(OllamaChatResponse.class)
                .map(item -> item.getMessage().getContent());
    }
}
